{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 1. 导入补偿飞行数据， 做低通， 构造训练数据， 再做带通/高通，生成最终训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from process_lib import pd_lowpass, pd_highpass, pd_bandpass, construct_features, calc_model, calc_compensation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "float_format = '%6.6e'\n",
    "\n",
    "# setup input filename\n",
    "filename = '4colsY2new.txt'\n",
    "# filename = '4cols.txt'\n",
    "# filename = '4cols2.txt'\n",
    "# filename = '4colsY2.txt'\n",
    "# filename = '4colUAV.txt'\n",
    "\n",
    "# setup output directory\n",
    "basename = os.path.splitext(filename)[0]\n",
    "input_dir = './input/'\n",
    "output_dir = './output/'\n",
    "intermediate_dir = './output/intermediate/'\n",
    "if not os.path.exists(output_dir):\n",
    "  # output_dir = os.path.join(os.path.curdir, 'output')\n",
    "  # intermediate_dir = os.path.join(output_dir, 'intermediate')\n",
    "  os.mkdir(output_dir)\n",
    "  os.mkdir(intermediate_dir)\n",
    "\n",
    "# input_file = os.path.join(input_dir, filename)\n",
    "input_file = input_dir + filename\n",
    "\n",
    "# setup filter options\n",
    "#_USING_HIGHPASS = False\n",
    "_USING_LOWPASS  = True\n",
    "_USING_SOS      = False\n",
    "print('_USING_LOWPASS:', _USING_LOWPASS)\n",
    "print('_USING_SOS:', _USING_SOS)\n",
    "\n",
    "# type = 'butter' #default\n",
    "# type = 'ellip'\n",
    "type = 'cheby1'\n",
    "# type = 'cheby2'\n",
    "order = 2;\n",
    "\n",
    "# for cheby1, cheby2, ellip, bessel\n",
    "rp = 1.\n",
    "rs = 60.\n",
    "\n",
    "# setup filter cutoff frequency\n",
    "wn1 = 0.1\n",
    "wn2 = [0.04, 0.1]\n",
    "wn3 = 0.04\n",
    "delta_t = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_base = input_dir + os.path.splitext(input_file)[0]\n",
    "lowpass_file = intermediate_dir + basename +'_lowpass.csv'\n",
    "features_file = intermediate_dir + basename + '_features.csv'\n",
    "features_bandpass_file =  intermediate_dir + basename + '_features_bandpass.csv'\n",
    "\n",
    "print('read input file: ', input_file)\n",
    "df = pd.read_csv(input_file, delim_whitespace=True)\n",
    "\n",
    "# print('save input to df0')\n",
    "# if filename == '4cols.txt':\n",
    "#     df = df.iloc[200:,:]\n",
    "# df0 = df\n",
    "\n",
    "# Is lowpass necessary?\n",
    "# perform lowpass filter\n",
    "if _USING_LOWPASS:\n",
    "    df.iloc[:,:4] = pd_lowpass(df.iloc[:,:4], wn1, order=order, type=type, rp=rp, rs=rs)\n",
    "    print('save lowpass file:', lowpass_file)\n",
    "    df.to_csv(lowpass_file, float_format=float_format, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. 构造16项训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct 16terms of features + label\n",
    "df_features = construct_features(df, delta_t = delta_t, base_name=basename, output_dir=intermediate_dir)\n",
    "print('save features_file: ', features_file)\n",
    "df_features.to_csv(features_file, float_format=float_format, sep='\\t', index=False)\n",
    "\n",
    "# perform band filter\n",
    "print('对(低通后的)输入三份量进行带通滤波')\n",
    "df_features_filtered = pd_highpass(df_features, wn3)\n",
    "print('save filered features: ', features_bandpass_file)\n",
    "df_features_filtered.to_csv(features_bandpass_file, float_format=float_format, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 2. 绘制训练数据(补偿飞行方向余弦及导数16项)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import rc\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "mpl.rcParams[\"font.size\"] = 4\n",
    "# plt.legend(fontsize = 9)\n",
    "# plt.xticks(fontsize = 10)\n",
    "# plt.yticks(fontsize = 10)\n",
    "\n",
    "fig, ax = plt.subplots(4,1)\n",
    "#rc[\"ax.labelsize\"] = 12\n",
    "fig.set_dpi(800)\n",
    "data1 = df_features.iloc[:,:3]\n",
    "#t1 = [i / 10 for i in range(len(data1.size))]\n",
    "t0 = [i for i in range(data1.shape[0])]\n",
    "t10 = [i / 10 for i in range(data1.shape[0])]\n",
    "#fft1 = scipy.fft.rfft(data1)\n",
    "ax[0].plot(t10, data1.iloc[:,0], linewidth=0.5, label=data1.columns[0])\n",
    "ax[0].plot(t10, data1.iloc[:,1], linewidth=0.5, label=data1.columns[1])\n",
    "ax[0].plot(t10, data1.iloc[:,2], linewidth=0.5, label=data1.columns[2])\n",
    "ax[0].legend(loc='upper right', fontsize = 3) \n",
    "\n",
    "data2 = df_features.iloc[:,3:8] #.to_numpy()\n",
    "ax[1].plot(t10, data2.iloc[:,0], linewidth=0.5, label=data2.columns[0])\n",
    "ax[1].plot(t10, data2.iloc[:,1], linewidth=0.5, label=data2.columns[1])\n",
    "ax[1].plot(t10, data2.iloc[:,2], linewidth=0.5, label=data2.columns[2])\n",
    "ax[1].plot(t10, data2.iloc[:,3], linewidth=0.5, label=data2.columns[3])\n",
    "ax[1].plot(t10, data2.iloc[:,4], linewidth=0.5, label=data2.columns[4])\n",
    "ax[1].legend(loc='upper right', fontsize = 3) \n",
    "\n",
    "data3 = df_features.iloc[:,8:16] #.to_numpy()\n",
    "ax[2].plot(t10, data3.iloc[:,0], linewidth=0.5, label=data3.columns[0])\n",
    "ax[2].plot(t10, data3.iloc[:,1], linewidth=0.5, label=data3.columns[1])\n",
    "ax[2].plot(t10, data3.iloc[:,2], linewidth=0.5, label=data3.columns[2])\n",
    "ax[2].plot(t10, data3.iloc[:,3], linewidth=0.5, label=data3.columns[3])\n",
    "ax[2].plot(t10, data3.iloc[:,4], linewidth=0.5, label=data3.columns[4])\n",
    "ax[2].plot(t10, data3.iloc[:,5], linewidth=0.5, label=data3.columns[5])\n",
    "ax[2].plot(t10, data3.iloc[:,6], linewidth=0.5, label=data3.columns[6])\n",
    "ax[2].plot(t10, data3.iloc[:,7], linewidth=0.5, label=data3.columns[7])\n",
    "ax[2].legend(loc='upper right', fontsize = 3) \n",
    "\n",
    "data4 = df_features.iloc[:,-1:] #.to_numpy()\n",
    "ax[3].plot(t0, data4.iloc[:,0], linewidth=0.5, label=data4.columns[0])\n",
    "ax[3].legend(loc='upper right', fontsize = 3) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3. 计算补偿数据的改善比（可用的回归算法：['OLS', 'RidgeCV', 'Ridge', 'PLSR', 'BayesianRidge', 'SM']） "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coef = None\n",
    "model_types = ['OLS', 'Ridge', 'PLSR', 'SM', 'BayesianRidge']\n",
    "STD_y = 0.0\n",
    "STD_Err = 0.0\n",
    "ir = None\n",
    "for model_type in model_types:\n",
    "    model, coef = calc_model(df_features_filtered, model_type)\n",
    "    #X_train = df_features_filtered.iloc[:,:16]\n",
    "    X_train = df_features_filtered.iloc[:,:-1]  # 除最后一列\n",
    "    y_train = df_features_filtered.iloc[:,-1].to_numpy()\n",
    "    y_pred  = model.predict(X_train)\n",
    "    STD_y   = np.std(y_train)\n",
    "    STD_Err = np.std(y_train - y_pred)\n",
    "    ir = STD_y, STD_Err, float(STD_y)/float(STD_Err)\n",
    "    print('model:%s\\n' %model_type, coef)\n",
    "    print('改善比: %2.2f %2.2f %2.2f\\n' %(STD_y, STD_Err, float(STD_y)/float(STD_Err)))\n",
    "    # if (model_type == 'SM'):\n",
    "    #     print(model.summary())    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. 计算butterworth参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import butter, filtfilt\n",
    "#order = order\n",
    "w1 = wn1 #0.1\n",
    "b1, a1 = butter(order, w1, btype='low')\n",
    "print('b1:', b1)\n",
    "print('a1:', a1)\n",
    "w2 = wn2 #[0.04, 0.1]\n",
    "b2, a2 = butter(order, w2, btype='band')\n",
    "print('b2:', b2)\n",
    "print('a2:',a2)\n",
    "sos1 = butter(order, w1, btype='low', output='sos')\n",
    "print('sos1:\\n',sos1)\n",
    "sos2 = butter(order, w2, btype='band', output='sos')\n",
    "print('sos2:\\n',sos2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5. 输出补偿系数、改善比到json文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "out_dict = {'name':model_types[-1], \n",
    "            'ir':ir, 'coef':coef.tolist(), \n",
    "            'wn1':wn1, 'wn2':wn2, 'wn3':wn3,'b1':b1.tolist(),\n",
    "            'a1':a1.tolist(),'b2':b2.tolist(),'a2':a2.tolist(),\n",
    "            'sos1':sos1.tolist(), 'sos2':sos2.tolist(),\n",
    "            '_USING_SOS':_USING_SOS, 'LowPass':_USING_LOWPASS,\n",
    "            'order':order, 'type':type, 'rp':rp, 'rs':rs, 'dt':delta_t}\n",
    "\n",
    "#json_file = output_dir + basename + '_compensation.json'\n",
    "json_file = output_dir + basename + '_compensation' + '%s' % order + '.json'\n",
    "if _USING_LOWPASS != True:\n",
    "  json_file = output_dir + basename + '_compensation_nolp.json'\n",
    "  json_file = output_dir + basename + '_compensation' + '%s' % order + '_nolp.json'\n",
    "#print(\"output to json file: \", json_file)\n",
    "#print(out_dict)\n",
    "json_object = json.dumps(out_dict, indent=4)\n",
    "print(json_object)\n",
    "with open(json_file, 'w') as ofile:\n",
    "  json.dump(out_dict, ofile)\n",
    "  # json.dump(json_object, ofile)\n",
    "print(\"Write to json file: \", json_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 6. 运行补偿(实际计算)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# df = pd.read_csv('y12实测数据2110.txt', delim_whitespace=True)\n",
    "# df_op = df.iloc[:,1:9]\n",
    "# df_op.to_csv('7colsY2_op2110.txt', float_format=' %6.5f', sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = '7colsY2_op9999.txt'\n",
    "# filename = '7colsY2_op2110.txt'\n",
    "# filename = '7colUAV_op1010.txt'\n",
    "basename = os.path.splitext(filename)[0]\n",
    "input_file = input_dir + filename\n",
    "input_base = input_dir + basename\n",
    "print('Read operation data from: ', input_file)\n",
    "df = pd.read_csv(input_file, delim_whitespace=True)\n",
    "\n",
    "if _USING_LOWPASS:\n",
    "  print(\"Do lowpass for 7 columns of input.\")\n",
    "  df = pd_lowpass(df, wn1, order=order, type=type, rp=rp, rs=rs)\n",
    "else:\n",
    "  print(\"Don't Do lowpass for 7 columns of input.\")\n",
    "\n",
    "df['未补磁低通1'] = df.iloc[:,3]\n",
    "\n",
    "print('构造16+1训练数据')\n",
    "df_features = construct_features(df, delta_t = delta_t, base_name=basename, output_dir=intermediate_dir)\n",
    "\n",
    "train_components = df_features\n",
    "df_features_file = intermediate_dir + basename + '_features.csv'\n",
    "print('Write features to: ', df_features_file)\n",
    "df_features.to_csv(df_features_file)\n",
    "\n",
    "df_features = pd_bandpass(df_features, wn2, order=order, type=type, rp=rp, rs=rs)\n",
    "\n",
    "df_features_bandpass_file = intermediate_dir + basename + '_features_bandpass.csv'\n",
    "print('Write filtered features to: ', df_features_bandpass_file)\n",
    "df_features.to_csv(df_features_bandpass_file)\n",
    "\n",
    "# prepair output columns (tobe apeending to origin file).\n",
    "# df['未补磁低通1'] = df.iloc[:,3]\n",
    "df['未补磁高通1'] = df_features.iloc[:,-1]\n",
    "compansation = calc_compensation(df_features, coef)\n",
    "df['补偿值1'] = compansation\n",
    "df['已补磁高通1'] = df['未补磁高通1'] - compansation\n",
    "df['已补磁总场1'] = df['未补磁总场'] - compansation\n",
    "\n",
    "compensated_file = output_dir + basename + '_final'  + '%s' % order + '.txt'\n",
    "if _USING_LOWPASS != True:\n",
    "  compensated_file = output_dir + basename + '_final_nolp.txt'\n",
    "print('Write to final compensated file: ', compensated_file)\n",
    "df.to_csv(compensated_file, float_format=float_format, sep='\\t', index=False, encoding='gbk')\n",
    "df5 = df.iloc[:,-5:]\n",
    "df5_file = intermediate_dir + basename + '_output'  + '%s' % order + '.txt'\n",
    "print('Write to final5 output file: ', df5_file)\n",
    "df5.to_csv(df5_file, float_format=float_format, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "\n",
    "# mpl.rcParams[\"font.size\"] = 7\n",
    "\n",
    "plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 显示中文\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "warnings.filterwarnings('ignore')\n",
    "t1 = [i / 10 for i in range(len(compansation))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig2 = plt.figure(num=1, figsize=(15, 8),dpi=800)     #开启一个窗口，同时设置大小，分辨率\n",
    "# ax21 = fig2.add_subplot(2,1,1)\n",
    "# ax22 = fig2.add_subplot(2,1,2)\n",
    "\n",
    "# #设置子图的基本元素\n",
    "# ax21.set_title('补偿对比图') \n",
    "# #plt.axis([0,800,-2,2])\n",
    "# # ax1.set_xlim(0,800)\n",
    "# # ax1.set_ylim(-2,2) \n",
    "# # ax2.set_xlim(0,800) \n",
    "# # ax2.set_ylim(-2,2)\n",
    "# # ax3.set_xlim(0,800) \n",
    "# # ax3.set_ylim(-1,1)\n",
    "# ax22.set_xlabel('time(s)')\n",
    "# plot1=ax21.plot(t1, df_op0.iloc[:,3], 'bo', linewidth=0.5, label=df_op0.columns[3])\n",
    "# plot1=ax21.plot(t1, df_op0.iloc[:,11], 'r-', linewidth=0.5, label=df_op0.columns[11])\n",
    "# ax21.legend(loc='upper right') \n",
    "# plot2=ax22.plot(t1, df_op0.iloc[:,11], linewidth=0.8, label=df_op0.columns[11])\n",
    "# # plot2=ax22.plot(t1, df_op0.iloc[:,12], linewidth=0.6, label=df_op0.columns[12])\n",
    "# # plot2=ax22.plot(t1, df_op0.iloc[:,13], linewidth=0.4, label=df_op0.columns[13])\n",
    "# ax22.legend(loc='upper right') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 6. 运行补偿(16 features plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3 = plt.figure(num=1, figsize=(15, 8),dpi=800)\n",
    "ax41 = fig3.add_subplot(4,1,1)\n",
    "ax42 = fig3.add_subplot(4,1,2)\n",
    "ax43 = fig3.add_subplot(4,1,3)\n",
    "ax44 = fig3.add_subplot(4,1,4)\n",
    "\n",
    "data1 = df_features.iloc[:,:3] #.to_numpy()\n",
    "#fft1 = scipy.fft.rfft(data1)\n",
    "#ax41.plot(data1, linewidth=0.5)\n",
    "ax41.plot(t1, data1.iloc[:,0], linewidth=0.5, label=data1.columns[0])\n",
    "ax41.plot(t1, data1.iloc[:,1], linewidth=0.5, label=data1.columns[1])\n",
    "ax41.plot(t1, data1.iloc[:,2], linewidth=0.5, label=data1.columns[2])\n",
    "ax41.legend(loc='upper right') \n",
    "\n",
    "data2 = df_features.iloc[:,3:8] #.to_numpy()\n",
    "#ax42.plot(data2, linewidth=0.5)\n",
    "ax42.plot(t1, data2.iloc[:,0], linewidth=0.5, label=data2.columns[0])\n",
    "ax42.plot(t1, data2.iloc[:,1], linewidth=0.5, label=data2.columns[1])\n",
    "ax42.plot(t1, data2.iloc[:,2], linewidth=0.5, label=data2.columns[2])\n",
    "ax42.plot(t1, data2.iloc[:,3], linewidth=0.5, label=data2.columns[3])\n",
    "ax42.plot(t1, data2.iloc[:,4], linewidth=0.5, label=data2.columns[4])\n",
    "ax42.legend(loc='upper right') \n",
    "\n",
    "data3 = df_features.iloc[:,8:16] #.to_numpy()\n",
    "#ax43.plot(data3, linewidth=0.5)\n",
    "ax43.plot(t1, data3.iloc[:,0], linewidth=0.5, label=data3.columns[0])\n",
    "ax43.plot(t1, data3.iloc[:,1], linewidth=0.5, label=data3.columns[1])\n",
    "ax43.plot(t1, data3.iloc[:,2], linewidth=0.5, label=data3.columns[2])\n",
    "ax43.plot(t1, data3.iloc[:,3], linewidth=0.5, label=data3.columns[3])\n",
    "ax43.plot(t1, data3.iloc[:,4], linewidth=0.5, label=data3.columns[4])\n",
    "ax43.plot(t1, data3.iloc[:,5], linewidth=0.5, label=data3.columns[5])\n",
    "ax43.plot(t1, data3.iloc[:,6], linewidth=0.5, label=data3.columns[6])\n",
    "ax43.plot(t1, data3.iloc[:,7], linewidth=0.5, label=data3.columns[7])\n",
    "ax43.legend(loc='upper right') \n",
    "\n",
    "data4 = df.iloc[:,11:14] #.to_numpy()\n",
    "ax44.plot(t1, data4.iloc[:,0], linewidth=0.5, label=data4.columns[0])\n",
    "# ax44.plot(t1, data4.iloc[:,1], linewidth=0.5, label=data4.columns[1])\n",
    "# ax44.plot(t1, data4.iloc[:,2], linewidth=0.5, label=data4.columns[2])\n",
    "ax44.legend(loc='upper right') \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.fft import fft, ifft, rfft, irfft\n",
    "\n",
    "x3 = df_features.iloc[:,8].to_numpy() # COSX*COSX_D\n",
    "Fs = 10\n",
    "N = len(x3)\n",
    "dt = 1.0/Fs\n",
    "t1 = np.arange(N)*dt\n",
    "t = np.linspace(0, (N-1)*dt, N)\n",
    "# 幅值谱计算\n",
    "# y3=np.abs(fft(x3))\n",
    "y3=np.abs(fft(x3))\n",
    "#f3=np.arange(int(N/2))*Fs/N;\n",
    "f3 = np.linspace(0., 0.1, int(N/2))\n",
    "# 绘制赋值谱\n",
    "fig4, ax=plt.subplots(nrows=2, ncols=1);\n",
    "fig4.dpi = 800\n",
    "ax[0].set_title('幅值谱')\n",
    "ax[0].plot(t, x3, linewidth=0.5, label=df_features.columns[8]);\n",
    "ax[0].legend(loc='upper right')\n",
    "ax[1].plot(f3, 2*y3[:int(N/2)]/N, linewidth=0.5, label='ABS(频谱)');\n",
    "ax[1].legend(loc='upper right')\n",
    "#plt.show();"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
